import os
import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from pydicom import Dataset
from pydicom.uid import generate_uid
import uuid
import json
import asyncio
from pathlib import Path
import zipfile
import io

app = FastAPI()

# 创建必要的目录
UPLOAD_DIR = Path("uploads")
DICOM_OUTPUT_DIR = Path("dicom_output")
UPLOAD_DIR.mkdir(exist_ok=True)
DICOM_OUTPUT_DIR.mkdir(exist_ok=True)

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")

# 存储转换进度
conversion_progress = {}

@app.get("/")
async def read_root():
    return FileResponse('static/index.html')

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    if not file.filename.endswith('.avi'):
        raise HTTPException(status_code=400, detail="只支持AVI格式文件")
    
    # 生成唯一任务ID
    task_id = str(uuid.uuid4())
    
    # 保存上传的文件
    file_path = UPLOAD_DIR / f"{task_id}_{file.filename}"
    with open(file_path, "wb") as buffer:
        content = await file.read()
        buffer.write(content)
    
    # 初始化进度
    conversion_progress[task_id] = {
        "status": "processing",
        "progress": 0,
        "total_frames": 0,
        "current_frame": 0,
        "dicom_files": []
    }
    
    # 异步开始转换
    asyncio.create_task(convert_avi_to_dicom(task_id, file_path, file.filename))
    
    return {"task_id": task_id, "filename": file.filename}

async def convert_avi_to_dicom(task_id: str, avi_path: Path, filename: str):
    try:
        # 打开视频文件
        cap = cv2.VideoCapture(str(avi_path))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        conversion_progress[task_id]["total_frames"] = total_frames
        
        # 为每个任务创建单独的输出目录
        task_output_dir = DICOM_OUTPUT_DIR / task_id
        task_output_dir.mkdir(exist_ok=True)
        
        dicom_files = []
        
        for frame_num in range(total_frames):
            ret, frame = cap.read()
            if not ret:
                break
            
            # 转换颜色空间（OpenCV使用BGR，DICOM需要RGB）
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # 创建DICOM文件
            ds = Dataset()
            
            # 设置基本DICOM属性
            ds.PatientName = f"Patient_{task_id}"
            ds.PatientID = task_id
            ds.StudyInstanceUID = generate_uid()
            ds.SeriesInstanceUID = generate_uid()
            ds.SOPInstanceUID = generate_uid()
            ds.SOPClassUID = "1.2.840.10008.5.1.4.1.1.2"  # CT Image Storage
            
            # 设置图像属性
            ds.Rows, ds.Columns, _ = frame_rgb.shape
            ds.SamplesPerPixel = 3
            ds.PhotometricInterpretation = "RGB"
            ds.BitsAllocated = 8
            ds.BitsStored = 8
            ds.HighBit = 7
            ds.PixelRepresentation = 0
            ds.PlanarConfiguration = 0
            
            # 设置像素数据
            ds.PixelData = frame_rgb.tobytes()
            
            # 设置帧信息
            ds.InstanceNumber = frame_num + 1
            
            # 保存DICOM文件
            dicom_filename = f"frame_{frame_num+1:04d}.dcm"
            dicom_path = task_output_dir / dicom_filename
            ds.save_as(str(dicom_path))
            
            dicom_files.append(dicom_filename)
            
            # 更新进度
            conversion_progress[task_id]["current_frame"] = frame_num + 1
            conversion_progress[task_id]["progress"] = int((frame_num + 1) / total_frames * 100)
            conversion_progress[task_id]["dicom_files"] = dicom_files
            
            # 小延迟避免阻塞
            await asyncio.sleep(0.01)
        
        cap.release()
        
        # 标记完成
        conversion_progress[task_id]["status"] = "completed"
        
        # 清理原始AVI文件
        avi_path.unlink(missing_ok=True)
        
    except Exception as e:
        conversion_progress[task_id]["status"] = "error"
        conversion_progress[task_id]["error"] = str(e)

@app.get("/progress/{task_id}")
async def get_progress(task_id: str):
    if task_id not in conversion_progress:
        raise HTTPException(status_code=404, detail="任务不存在")
    return conversion_progress[task_id]

@app.get("/download/{task_id}/{filename}")
async def download_dicom(task_id: str, filename: str):
    file_path = DICOM_OUTPUT_DIR / task_id / filename
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="文件不存在")
    return FileResponse(str(file_path), filename=filename)

@app.get("/download-zip/{task_id}")
async def download_zip(task_id: str):
    task_dir = DICOM_OUTPUT_DIR / task_id
    if not task_dir.exists():
        raise HTTPException(status_code=404, detail="任务不存在")
    
    # 创建ZIP文件
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        for dcm_file in task_dir.glob("*.dcm"):
            zip_file.write(str(dcm_file), dcm_file.name)
    
    zip_buffer.seek(0)
    
    return StreamingResponse(
        io.BytesIO(zip_buffer.getvalue()),
        media_type="application/zip",
        headers={"Content-Disposition": f"attachment; filename=dicom_{task_id}.zip"}
    )

@app.get("/list-files/{task_id}")
async def list_files(task_id: str):
    task_dir = DICOM_OUTPUT_DIR / task_id
    if not task_dir.exists():
        return {"files": []}
    
    files = [f.name for f in task_dir.glob("*.dcm")]
    return {"files": files}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)